package com.lwc.backendsinglestructure.config;

import com.lwc.backendsinglestructure.common.utils.TraceIdUtil;
import com.lwc.backendsinglestructure.constants.LogConstant;
import com.lwc.backendsinglestructure.utils.ThreadLocalUtil;
import org.slf4j.MDC;
import org.springframework.core.task.TaskDecorator;

import java.util.Map;

/**
 * MDC 任务装饰器
 * 作用： 用于异步任务中，将父线程的MDC上下文传递给子线程。
 * 场景： 在异步任务执行中，传递日志链路ID追踪日志
 */
public class MdcTaskDecorator implements TaskDecorator {

    /**
     * @param runnable 原始任务
     * @return 装饰后的任务
     */
    @Override
    public Runnable decorate(Runnable runnable) {
        // 获取父线程MCD上下文
        Map<String, String> parentContextMap = MDC.getCopyOfContextMap();

        // 获取Trace_id 链路ID
        String TRACE_ID = parentContextMap != null ? parentContextMap.get(LogConstant.TRACE_ID) : TraceIdUtil.getTraceId();

        return () -> {
            try {
                // 设置MDC
                if (parentContextMap == null) {
                    MDC.put(LogConstant.TRACE_ID, TRACE_ID);
                }else{
                    // 将父线程的上下文设置到子线程中
                    MDC.setContextMap(parentContextMap);
                }
                // 执行任务
                runnable.run();
            } finally {
                // 清除MDC上下文，避免内存泄露
                MDC.clear();
            }
        };
    }
}